import torch
import torch.nn as nn
from collections import Counter
# from torch.fft import fft, fftfreq
class RevIN(nn.Module):
    def __init__(self, num_features: int, eps=1e-5, affine=True, subtract_last=False):
        """
        :param num_features: the number of features or channels
        :param eps: a value added for numerical stability
        :param affine: if True, RevIN has learnable affine parameters
        """
        super(RevIN, self).__init__()
        self.num_features = num_features
        self.eps = eps
        self.affine = affine
        self.subtract_last = subtract_last
        if self.affine:
            self._init_params()

    def forward(self, x, mode:str):
        if mode == 'norm':
            self._get_statistics(x)
            x = self._normalize(x)
        elif mode == 'denorm':
            x = self._denormalize(x)
        else: raise NotImplementedError
        return x

    def _init_params(self):
        # initialize RevIN params: (C,)
        self.affine_weight = nn.Parameter(torch.ones(self.num_features))
        self.affine_bias = nn.Parameter(torch.zeros(self.num_features))

    def _get_statistics(self, x):
        dim2reduce = tuple(range(1, x.ndim-1))
        if self.subtract_last:
            self.last = x[:,-1,:].unsqueeze(1)
        else:
            self.mean = torch.mean(x, dim=dim2reduce, keepdim=True).detach()
        self.stdev = torch.sqrt(torch.var(x, dim=dim2reduce, keepdim=True, unbiased=False) + self.eps).detach()

    def _normalize(self, x):
        if self.subtract_last:
            x = x - self.last
        else:
            x = x - self.mean
        x = x / self.stdev
        if self.affine:
            x = x * self.affine_weight
            x = x + self.affine_bias
        return x

    def _denormalize(self, x):
        if self.affine:
            x = x - self.affine_bias
            x = x / (self.affine_weight + self.eps*self.eps)
        x = x * self.stdev
        if self.subtract_last:
            x = x + self.last
        else:
            x = x + self.mean
        return x

# class cal_patch_len(nn.Module):
#
#     def calculate_patch_length(x, fs=100, num_peaks=1, default_patch_len=12):
#         """
#         计算基于信号频率成分的 patch 长度。
#
#         参数:
#             x (torch.Tensor): 输入张量，形状为 (batch_size, seq_len)。
#             fs (float): 采样频率。
#             num_peaks (int): 关注的显著频率成分数量，默认为 1。
#             default_patch_len (int): 如果没有有效周期值时的默认 patch 长度，默认为 12。
#         返回:
#             int: 计算得到的 patch 长度。
#         """
#         # 计算 FFT
#         fft_values = torch.fft.fft(x)
#         fft_magnitude = torch.abs(fft_values)
#         seq_len = x.shape[1]
#         fft_frequency = torch.fft.fftfreq(seq_len, d=1 / fs, device=x.device)
#
#         # 由于对称性，只取一半的频率范围
#         half_length = seq_len // 2 + 1
#         fft_magnitude = fft_magnitude[:, :half_length]
#         fft_frequency = fft_frequency[:half_length]
#
#         # 识别显著的频率和周期
#         peak_indices = torch.topk(fft_magnitude, k=num_peaks, dim=1).indices
#         peak_frequencies = torch.gather(fft_frequency.unsqueeze(0).expand(x.shape[0], -1), 1, peak_indices)
#
#         # 将所有周期转换为整数（假设周期是整数）
#         peak_periods_int = torch.round(1 / peak_frequencies).to(torch.int32)
#
#         # 过滤掉无效的周期值（例如，周期为零或负值）
#         valid_periods = peak_periods_int[peak_periods_int > 0].cpu().numpy().flatten()
#
#         # 如果没有有效的周期值，设置默认值
#         if len(valid_periods) == 0:
#             patch_len = default_patch_len
#             print(f"No valid periods found. Setting default patch length to {default_patch_len}.")
#         else:
#             # 统计每个周期的出现次数
#             period_counts = Counter(valid_periods)
#
#             # 选择出现次数最多的周期
#             most_common_period = period_counts.most_common(1)[0][0]
#
#             # 如果所有周期都只出现一次，则选择最小的周期
#             if period_counts[most_common_period] == 1:
#                 patch_len = min(valid_periods)
#                 if patch_len == 1:
#                     patch_len = default_patch_len
#                 print("All periods appear only once. Setting patch length to the smallest period.")
#             else:
#                 patch_len = most_common_period
#
#         print(f"Patch length: {patch_len}")
#         return patch_len

